/*
* Copyright (C) Microsoft. All rights reserved.  
*
* Licensed under the Apache License, Version 2.0 (the "License"); you may not 
* use this file except in compliance with the License.  You may obtain a copy 
* of the License at http://www.apache.org/licenses/LICENSE-2.0  
* 
* THIS CODE IS PROVIDED *AS IS* BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
* KIND, EITHER EXPRESS OR IMPLIED, INCLUDING WITHOUT LIMITATION ANY IMPLIED 
* WARRANTIES OR CONDITIONS OF TITLE, FITNESS FOR A PARTICULAR PURPOSE, 
* MERCHANTABLITY OR NON-INFRINGEMENT. 
*
* See the Apache License, Version 2.0 for specific language governing 
* permissions and limitations under the License.
*
*/
#pragma once

#include <baja/graphics/dx11.1/details/common.hpp>
#include <baja/graphics/dx11.1/graphics_fwd.hpp>
#include <baja/graphics/dx11.1/details/enum_mappings.hpp>
#include <baja/graphics/dx11.1/details/util.hpp>
#include <baja/stream/stream.hpp>

namespace baja { namespace graphics { namespace dx11_1 {

namespace d = baja::graphics::dx11_1::details;

// this is a very simple way of defining "effects"
// since fx is not supported in metro, and the original graphics engine was strongly tied into fx, I came up
// with this class that is a fairly crappy, yet effecive way of defining an effect consisting of :
// - a pixel shader
// - a vertex shader
// - a list of constant buffers 
// - a list of samplers
// - a list of textures
// - a list of sampler states
// 
// these are all defined in code in a "effect2_description" structure which is used to instantiate the effect2 here
// some default constant buffers, textures and the single texture sampler are defined in
// "/baja/graphics/dx11.1/hlsl/common_params.hlsli"
// this can be included in the hlsl file to use these.  make sure constant buffer slots begin in 5 to not collide with these
class effect2 : public ieffect2, public b::iunknown_provider, public std::enable_shared_from_this<effect2>
{
public:
    // constructor
    effect2(
        const std::shared_ptr<iengine>& graphics,
        const std::shared_ptr<itexture>& nullTexture
        ) : _graphics(graphics), _nullTexture(nullTexture)
    {
    }
    static std::shared_ptr<effect2> create(
        const std::shared_ptr<iengine>& graphics,
         const std::shared_ptr<effect2_description>& desc,
         const std::shared_ptr<itexture>& nullTexture
        )
    {
        std::shared_ptr<effect2> e = std::make_shared<effect2>(graphics, nullTexture);
        e->initialize(desc);
        return e;
    }
    
    // b::iunknown_provider
    size_t getInterfaceCount() { return 0; }

    WRL::ComPtr<IUnknown> getInterface(size_t index)
    {
        return nullptr;
    }   

    void setTexture(const wchar* name, const std::shared_ptr<itexture>& texture, const uint32 context = 0)
    {
        std::shared_ptr<itexture> t = texture;
        if (!t)
        {
            t = _nullTexture;
        }

        std::shared_ptr<b::iunknown_provider> w = std::dynamic_pointer_cast<b::iunknown_provider>(_graphics.lock());
        WRL::ComPtr<IUnknown> punk = w->getInterface((context == 0) ? d::graphics_interfaces::d3d_device_context : d::graphics_interfaces::d3d_device_context_deferred);
        WRL::ComPtr<ID3D11DeviceContext> deviceContext;
        punk.As(&deviceContext);

        WRL::ComPtr<ID3D11ShaderResourceView> shaderResource;
        w = std::dynamic_pointer_cast<b::iunknown_provider>(t);
        punk = w->getInterface(d::texture_interfaces::shader_resource);  
        punk.As(&shaderResource);

        auto it = _textureMap.find(name);
        if (it == _textureMap.end())
        {
            return;
        }
        std::shared_ptr<texture_description> holder = it->second;

        deviceContext->PSSetShaderResources(holder->slot, 1, shaderResource.GetAddressOf());
    }

    void setConstantBuffer(const wchar* name, const void* buffer, const uint32 context = 0)
    {
        auto it = _constantBufferMap.find(name);
        BAJA_VERIFY(it != _constantBufferMap.end());
        std::shared_ptr<constant_buffer_holder> holder = it->second;

        std::shared_ptr<b::iunknown_provider> w = std::dynamic_pointer_cast<b::iunknown_provider>(_graphics.lock());
        WRL::ComPtr<IUnknown> punk = w->getInterface((context == 0) ? d::graphics_interfaces::d3d_device_context : d::graphics_interfaces::d3d_device_context_deferred);
        WRL::ComPtr<ID3D11DeviceContext> deviceContext;
        punk.As(&deviceContext);

        D3D11_MAPPED_SUBRESOURCE mappedResource;
        win32_result((HRESULT)deviceContext->Map(holder->cb.Get(), 0, D3D11_MAP_WRITE_DISCARD, 0, &mappedResource));
        CopyMemory(mappedResource.pData, buffer, holder->desc.byteSize);
        deviceContext->Unmap(holder->cb.Get(), 0);

        deviceContext->VSSetConstantBuffers(
            holder->desc.slot,
            1,
            holder->cb.GetAddressOf()
            );
        deviceContext->PSSetConstantBuffers(
            holder->desc.slot,
            1,
            holder->cb.GetAddressOf()
            );
    }

    void setBlendState(const wchar* name, const float32* blendFactor, uint32 sampleMask, const uint32 context = 0)
    {
        auto it = _blendStateMap.find(name);
        BAJA_VERIFY(it != _blendStateMap.end());
        std::shared_ptr<blend_state_holder> holder = it->second;

        std::shared_ptr<b::iunknown_provider> w = std::dynamic_pointer_cast<b::iunknown_provider>(_graphics.lock());
        WRL::ComPtr<IUnknown> punk = w->getInterface((context == 0) ? d::graphics_interfaces::d3d_device_context : d::graphics_interfaces::d3d_device_context_deferred);
        WRL::ComPtr<ID3D11DeviceContext> deviceContext;
        punk.As(&deviceContext);

        deviceContext->OMSetBlendState(holder->bs.Get(), blendFactor, sampleMask);
    }

    // NOTE: we only support one sampler in slot 0 right now
    void setSamplerState(const wchar* name, const uint32 context = 0)
    {
        auto it = _samplerStateMap.find(name);
        BAJA_VERIFY(it != _samplerStateMap.end());
        std::shared_ptr<sampler_state_holder> holder = it->second;

        std::shared_ptr<b::iunknown_provider> w = std::dynamic_pointer_cast<b::iunknown_provider>(_graphics.lock());
        WRL::ComPtr<IUnknown> punk = w->getInterface((context == 0) ? d::graphics_interfaces::d3d_device_context : d::graphics_interfaces::d3d_device_context_deferred);
        WRL::ComPtr<ID3D11DeviceContext> deviceContext;
        punk.As(&deviceContext);

        deviceContext->PSSetSamplers(0, 1, holder->ss.GetAddressOf());
    }

    size_t pixelShaderSize() { return _pixelShaderFileStream->getBytes().size(); }
    const void* pixelShaderBytes() { return _pixelShaderFileStream->getBytes().data(); }
    size_t vertexShaderSize() { return _vertexShaderFileStream->getBytes().size(); }
    const void* vertexShaderBytes() { return _vertexShaderFileStream->getBytes().data(); }

    void begin(const std::shared_ptr<iengine>& graphics, const uint32 context = 0)
    {
        win32_result wr;
        
        // cache current states to restore once effect is done
        std::shared_ptr<engine> graphicsImpl = std::dynamic_pointer_cast<engine>(graphics);
        graphicsImpl->getCurrentBlendState(
            &_blendStateCache.blendStateDesc,
            _blendStateCache.blendFactor,
            &_blendStateCache.sampleMask,
            context
            );

        graphicsImpl->getCurrentDepthStencilState(
            &_depthStencilStateCache.depthStencilStateDesc,
            &_depthStencilStateCache.stencilRef,
            context
            );

        graphicsImpl->getCurrentRasterizerState(
            &_rasterizerStateCache.rasterizerStateDesc,
            context
            );

        // get the device context
        WRL::ComPtr<ID3D11DeviceContext> deviceContext;
        std::shared_ptr<b::iunknown_provider> w = std::dynamic_pointer_cast<b::iunknown_provider>(_graphics.lock());
        WRL::ComPtr<IUnknown> punk = w->getInterface((context == 0) ? d::graphics_interfaces::d3d_device_context : d::graphics_interfaces::d3d_device_context_deferred);
        punk.As(&deviceContext);

        // Set the vertex shader
        deviceContext->VSSetShader(
            _vertexShader.Get(),
            nullptr,
            0
            );

        // Set the pixel shader
        graphicsImpl->setPixelShader(_pixelShader, context);

        graphics->updateEffectTextures(std::static_pointer_cast<g::ieffect2>(shared_from_this()), context);
    }
        
    void end(const std::shared_ptr<iengine>& graphics, const uint32 context = 0)
    {
        win32_result wr;
        
        std::shared_ptr<engine> graphicsImpl = std::dynamic_pointer_cast<engine>(graphics);

        // restore cached state
        graphicsImpl->setBlendState(
            _blendStateCache.blendStateDesc,
            _blendStateCache.blendFactor,
            _blendStateCache.sampleMask,
            context);

        graphicsImpl->setDepthStencilState(
            _depthStencilStateCache.depthStencilStateDesc,
            _depthStencilStateCache.stencilRef,
            context);

        graphicsImpl->setRasterizerState(_rasterizerStateCache.rasterizerStateDesc, context);

        setTexture(L"texture1", nullptr, context);
        setTexture(L"texture2", nullptr, context);
        setTexture(L"texture3", nullptr, context);
        setTexture(L"texture4", nullptr, context);
        setTexture(L"texture5", nullptr, context);
        setTexture(L"texture6", nullptr, context);
        setTexture(L"texture7", nullptr, context);
        setTexture(L"texture8", nullptr, context);

        // unbind stream 0
        graphics->setStreamSource(
            0, 
            std::shared_ptr<ivertex_buffer>(),
            0,
            sizeof(vertex_lit_colored_textured),
            context
            );
    }

private:
    void initialize(const std::shared_ptr<effect2_description>& desc)
    {
        win32_result wr;

        std::shared_ptr<b::iunknown_provider> w = std::dynamic_pointer_cast<b::iunknown_provider>(_graphics.lock());
        WRL::ComPtr<IUnknown> punk = w->getInterface(d::graphics_interfaces::d3d_device);
        WRL::ComPtr<ID3D11Device> device;
        punk.As(&device);

        // Create the pixel shader
        _pixelShaderFileStream = createMemoryStreamFromFile(desc->pixelShaderFile);
        wr = device->CreatePixelShader(
            _pixelShaderFileStream->getBytes().data(),
            _pixelShaderFileStream->getBytes().size(),
            nullptr,
            _pixelShader.GetAddressOf()
            );

        // Create the vertex shader
        _vertexShaderFileStream = createMemoryStreamFromFile(desc->vertexShaderFile);
        wr = device->CreateVertexShader(
            _vertexShaderFileStream->getBytes().data(),
            _vertexShaderFileStream->getBytes().size(),
            nullptr,
            _vertexShader.GetAddressOf()
            );

        // Set up the constant buffer map
        std::for_each(
            desc->constantBufferDesc.begin(),
            desc->constantBufferDesc.end(),
            [&](const constant_buffer_description& cbd)
        {
            WRL::ComPtr<ID3D11Buffer> cb;

            D3D11_BUFFER_DESC cd;
            ZeroMemory(&cd, sizeof(D3D11_BUFFER_DESC));
            cd.ByteWidth = cbd.byteSize;
            cd.Usage = D3D11_USAGE_DYNAMIC;
            cd.BindFlags = D3D11_BIND_CONSTANT_BUFFER;
            cd.CPUAccessFlags = D3D11_CPU_ACCESS_WRITE;
            cd.MiscFlags = 0;
            cd.StructureByteStride = 0;
    
            wr = device->CreateBuffer(
                &cd,
                nullptr,
                cb.GetAddressOf()
                );

            std::shared_ptr<constant_buffer_holder> holder = std::make_shared<constant_buffer_holder>(cbd);
            holder->cb = cb;
            _constantBufferMap[cbd.name] = holder;
        });

        // Set up the texture buffer map
        std::for_each(
            desc->textureDesc.begin(),
            desc->textureDesc.end(),
            [&](const texture_description& td)
        {
            std::shared_ptr<texture_description> holder = std::make_shared<texture_description>(td);
            _textureMap[td.name] = holder;
        });

        // Set up the blend state map
        std::for_each(
            desc->blendStateDesc.begin(),
            desc->blendStateDesc.end(),
            [&](const state_description& sd)
        {
            std::shared_ptr<blend_state_holder> holder = std::make_shared<blend_state_holder>(sd);

            wr = device->CreateBlendState(reinterpret_cast<const D3D11_BLEND_DESC *>(sd.bytes.data()), holder->bs.GetAddressOf());

            _blendStateMap[sd.name] = holder;
        });

        // Set up the sampler state map
        std::for_each(
            desc->samplerStateDesc.begin(),
            desc->samplerStateDesc.end(),
            [&](const state_description& sd)
        {
            std::shared_ptr<sampler_state_holder> holder = std::make_shared<sampler_state_holder>(sd);

            wr = device->CreateSamplerState(reinterpret_cast<const D3D11_SAMPLER_DESC *>(sd.bytes.data()), holder->ss.GetAddressOf());

            _samplerStateMap[sd.name] = holder;
        });
    }

private:
    std::shared_ptr<g::itexture> _nullTexture;

    struct blend_state_cache
    {
        D3D11_BLEND_DESC blendStateDesc;
        float32 blendFactor[4];
        uint32 sampleMask;
    };

    struct depth_stencil_state_cache
    {
        D3D11_DEPTH_STENCIL_DESC depthStencilStateDesc;
        uint32 stencilRef;
    };

    struct rasterizer_state_cache
    {
        D3D11_RASTERIZER_DESC rasterizerStateDesc;
    };

    blend_state_cache _blendStateCache;
    depth_stencil_state_cache _depthStencilStateCache;
    rasterizer_state_cache _rasterizerStateCache;

    std::weak_ptr<iengine> _graphics;
    std::shared_ptr<imemory_stream> _pixelShaderFileStream;
    std::shared_ptr<imemory_stream> _vertexShaderFileStream;
    WRL::ComPtr<ID3D11PixelShader> _pixelShader;
    WRL::ComPtr<ID3D11VertexShader> _vertexShader;

    struct constant_buffer_holder 
    { 
        constant_buffer_description desc; 
        WRL::ComPtr<ID3D11Buffer> cb; 

        constant_buffer_holder(const constant_buffer_description& desc_) : desc(desc_) {}
        constant_buffer_holder(const constant_buffer_holder& other) : desc(other.desc), cb(other.cb) {}
    };

    std::map<const std::wstring, std::shared_ptr<constant_buffer_holder>> _constantBufferMap;

    std::map<const std::wstring, std::shared_ptr<texture_description>> _textureMap;

    struct blend_state_holder
    { 
        state_description desc; 
        WRL::ComPtr<ID3D11BlendState> bs;

        blend_state_holder(const state_description& desc_) : desc(desc_) {}
        blend_state_holder(const blend_state_holder& other) : desc(other.desc), bs(other.bs) {}
    };

    std::map<const std::wstring, std::shared_ptr<blend_state_holder>> _blendStateMap;

    struct sampler_state_holder
    { 
        state_description desc; 
        WRL::ComPtr<ID3D11SamplerState> ss;

        sampler_state_holder(const state_description& desc_) : desc(desc_) {}
        sampler_state_holder(const sampler_state_holder& other) : desc(other.desc), ss(other.ss) {}
    };

    std::map<const std::wstring, std::shared_ptr<sampler_state_holder>> _samplerStateMap;
};

inline std::shared_ptr<ieffect2> engine::createEffect(const std::shared_ptr<effect2_description>& desc)
{ 
    std::shared_ptr<effect2> returnValue = effect2::create(shared_from_this(), desc, _nullTexture);

    return std::dynamic_pointer_cast<ieffect2>(returnValue);
}

}}} // namespace baja.graphics.d3d11_1
